import os
import yaml
import httpx
from typing import Optional, Dict, Any

# 读取 Ollama 地址配置，优先环境变量，其次 config.yaml，默认 http://localhost:11434
def get_ollama_base_url() -> str:
    env_url = os.getenv('OLLAMA_BASE_URL')
    if env_url:
        return env_url
    config_path = os.path.join(os.path.dirname(__file__), '../../config.yaml')
    if os.path.exists(config_path):
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
            if config and 'ollama_base_url' in config:
                return config['ollama_base_url']
    return 'http://localhost:11434'

OLLAMA_BASE_URL = get_ollama_base_url()
OLLAMA_MODEL = os.getenv('OLLAMA_MODEL', 'qwen3:8B')

class Qwen3Client:
    def __init__(self, base_url: Optional[str] = None, model: Optional[str] = None):
        self.base_url = base_url or OLLAMA_BASE_URL
        self.model = model or OLLAMA_MODEL
        self.api_url = f"{self.base_url}/api/generate"

    def _build_payload(self, prompt: str, **kwargs) -> Dict[str, Any]:
        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False
        }
        payload.update(kwargs)
        return payload

    def call(self, prompt: str, timeout: float = 5.0, **kwargs) -> str:
        """
        同步调用 Ollama Qwen3:8B
        """
        payload = self._build_payload(prompt, **kwargs)
        try:
            with httpx.Client(timeout=timeout) as client:
                resp = client.post(self.api_url, json=payload)
                if resp.status_code == 200:
                    data = resp.json()
                    return data.get('response', '')
                elif resp.status_code == 400:
                    raise ValueError(f"参数错误: {resp.text}")
                elif resp.status_code == 503:
                    raise RuntimeError("模型服务不可用")
                else:
                    raise RuntimeError(f"未知错误: {resp.status_code} {resp.text}")
        except httpx.RequestError as e:
            raise RuntimeError(f"请求失败: {e}")

    async def acall(self, prompt: str, timeout: float = 5.0, **kwargs) -> str:
        """
        异步调用 Ollama Qwen3:8B
        """
        payload = self._build_payload(prompt, **kwargs)
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                resp = await client.post(self.api_url, json=payload)
                if resp.status_code == 200:
                    data = resp.json()
                    return data.get('response', '')
                elif resp.status_code == 400:
                    raise ValueError(f"参数错误: {await resp.aread()}")
                elif resp.status_code == 503:
                    raise RuntimeError("模型服务不可用")
                else:
                    raise RuntimeError(f"未知错误: {resp.status_code} {await resp.aread()}")
        except httpx.RequestError as e:
            raise RuntimeError(f"请求失败: {e}") 